[DRAFT] Add native GQA attention support for ORT GenAI compatibility - #109
[DRAFT] Add native GQA attention support for ORT GenAI compatibility#109gramalingam wants to merge 1 commit into
Conversation
Add a new GQAAttention component that emits com.microsoft::GroupQueryAttention directly, and a gqa=True mode on CausalLMTask that produces models compatible with the onnxruntime-genai runtime. Key changes: - New GQAAttention component with GQAContext NamedTuple - DecoderLayer accepts pluggable attention_class parameter - CausalLMTask(gqa=True) builds graphs with: - No position_ids input (RoPE fused inside GQA via do_rotary=1) - seqlens_k/total_seq_len computed from attention_mask - cos_cache/sin_cache as graph initializers - com.microsoft opset import - CLI --gqa flag and build() module_kwargs parameter - 10 new GQA-specific tests in build_graph_test.py Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
There was a problem hiding this comment.
Pull request overview
This PR adds a first-class “native GQA” build path for causal language models to emit com.microsoft::GroupQueryAttention directly, targeting ONNX Runtime GenAI compatibility (fused RoPE + in-place KV cache). It threads a GQAContext through the model stack, adds an attention_class injection point in decoder layers, exposes a --gqa CLI flag, and introduces targeted graph-construction tests.
Changes:
- Add
GQAAttention/GQAContextcomponent emittingcom.microsoft::GroupQueryAttention. - Add
CausalLMTask(gqa=True)path that builds graphs withoutposition_ids, computesseqlens_k/total_seq_len, and registers cos/sin caches as initializers. - Add
attention_classplumbing +--gqaCLI flag andbuild(..., module_kwargs=...)support, plus new build-graph tests.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/build_graph_test.py | Adds a new TestBuildGQAGraph suite asserting expected I/O and ops for GQA graphs. |
| src/mobius/tasks/_causal_lm.py | Introduces gqa=True build mode (_build_gqa) and rotary cache discovery/registration. |
| src/mobius/models/base.py | Threads attention_class + gqa_context through base text/LM models and decoder invocation. |
| src/mobius/components/_gqa_attention.py | New component implementing GroupQueryAttention emission and the GQAContext container. |
| src/mobius/components/_decoder.py | Adds attention_class override and a GQA dispatch path in DecoderLayer. |
| src/mobius/components/init.py | Re-exports GQAAttention and GQAContext from the components package. |
| src/mobius/_builder.py | Adds module_kwargs passthrough to module constructors in build(). |
| src/mobius/main.py | Adds --gqa flag and wires it to CausalLMTask(gqa=True) + attention_class=GQAAttention. |
| # GQA mode: when attention_bias is a GQAContext, the attention | ||
| # component is GQAAttention and handles masking + RoPE internally. | ||
| if isinstance(attention_bias, GQAContext): | ||
| gqa_context = attention_bias | ||
| return self._forward_gqa( | ||
| op, | ||
| hidden_states, | ||
| gqa_context, | ||
| past_key_value, | ||
| ) |
There was a problem hiding this comment.
DecoderLayer.forward() routes to _forward_gqa() whenever attention_bias is a GQAContext, before checking self._post_norm. For post-norm layers, _forward_gqa() will crash (e.g., input_layernorm is not created when post_norm=True) and even if it didn’t, it would apply the wrong residual/norm ordering. Consider either (a) implementing a post-norm GQA path and dispatching based on self._post_norm, or (b) explicitly rejecting GQAContext on post-norm layers with a clear error so --gqa can’t silently build an invalid graph.
| def __init__( | ||
| self, | ||
| *, | ||
| static_cache: bool = False, | ||
| max_seq_len: int | None = None, | ||
| gqa: bool = False, | ||
| ): | ||
| self._static_cache = static_cache | ||
| self._max_seq_len = max_seq_len | ||
| self._gqa = gqa | ||
|
|
||
| def build( | ||
| self, | ||
| module: nn.Module, | ||
| config: ArchitectureConfig, | ||
| ) -> ModelPackage: | ||
| if self._gqa: | ||
| return self._build_gqa(module, config) | ||
|
|
There was a problem hiding this comment.
CausalLMTask allows static_cache=True and gqa=True to be set together via the Python API, but build() silently prioritizes _build_gqa() and ignores static_cache/max_seq_len. Add an explicit validation (in __init__ or build()) to reject the unsupported combination so callers don’t get an unexpected cache mode.
| elif args.gqa: | ||
| from mobius.components import GQAAttention | ||
|
|
||
| task = CausalLMTask(gqa=True) | ||
| module_kwargs = {"attention_class": GQAAttention} | ||
| trust_remote_code = args.trust_remote_code |
There was a problem hiding this comment.
--gqa unconditionally sets module_kwargs = {"attention_class": GQAAttention} and then instantiates module_class(config, **module_kwargs). Many registered model classes don’t accept an attention_class kwarg (e.g., models with custom __init__ signatures), so this will raise a TypeError with a confusing message. Consider detecting support (e.g., via inspect.signature / try/except TypeError) and failing with a clear CLI error listing supported model types, or only injecting attention_class for known compatible base models.
| dtype = resolve_dtype(dtype) | ||
| config = dataclasses.replace(config, dtype=dtype) | ||
|
|
||
| if task is None: | ||
| task = _default_task_for_model(model_type) | ||
|
|
||
| model_module = module_class(config) | ||
| model_module = module_class(config, **(module_kwargs or {})) | ||
| pkg = build_from_module(model_module, config, task) |
There was a problem hiding this comment.
build(..., module_kwargs=...) blindly forwards **module_kwargs into module_class(config, ...). If module_kwargs contains keys not accepted by a given model constructor (likely for many registry entries), this raises TypeError without context. Consider validating keys against the constructor signature or catching TypeError and re-raising a more actionable error that includes model_type and the unexpected kwargs.
| for layer, past_kv in zip(self.layers, past_kvs): | ||
| hidden_states, present_kv = layer( | ||
| op, | ||
| hidden_states=hidden_states, | ||
| attention_bias=gqa_context, | ||
| position_embeddings=None, | ||
| past_key_value=past_kv, | ||
| ) |
There was a problem hiding this comment.
This new GQA path passes attention_bias=gqa_context (a GQAContext) and position_embeddings=None into DecoderLayer.forward(), but the method is still annotated as attention_bias: ir.Value | None and position_embeddings: tuple. With the repo’s strict mypy settings, this should produce type errors. Update the DecoderLayer.forward (and any helpers) type annotations to accept GQAContext and None for position_embeddings (or introduce a dedicated gqa_context parameter instead of overloading attention_bias).
| ) | ||
|
|
||
|
|
||
| def _find_rotary_emb(module: nn.Module) -> nn.Module | None: |
There was a problem hiding this comment.
This doesn't look good (searching for a module by name and attribute). I wonder why this is necessary.
|
Superseded by #134 |
Note: This is a copilot-authored PR. I haven't yet checked it. Creating the PR for easy review and collaboration.
Add a new GQAAttention component that emits com.microsoft::GroupQueryAttention directly, and a gqa=True mode on CausalLMTask that produces models compatible with the onnxruntime-genai runtime.
Key changes: